import torch
from torch import nn, optim
from torch.nn import functional as F
import torchvision
from matplotlib import pyplot as plt

from utils import plot_image, plot_curve

# Define batch size and device
batch_size = 512
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# Load dataset with data augmentation
def load_data(batch_size):
    transform_train = torchvision.transforms.Compose([
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.RandomRotation(10),
        torchvision.transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
        torchvision.transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.1307,), (0.3081,))
    ])

    transform_test = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = torchvision.datasets.MNIST('mnist_data', train=True, download=True, transform=transform_train)
    test_dataset = torchvision.datasets.MNIST('mnist_data', train=False, download=True, transform=transform_test)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader


train_loader, test_loader = load_data(batch_size)


# Define the improved network structure
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


net = Net().to(device)
optimizer = optim.Adam(net.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()


# Train the network
def train(net, train_loader, optimizer, criterion, device, epochs=5):
    net.train()
    train_loss = []

    for epoch in range(epochs):
        for batch_idx, (x, y) in enumerate(train_loader):
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            output = net(x)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()

            train_loss.append(loss.item())

            if batch_idx % 10 == 0:
                print(f'Epoch {epoch + 1}/{epochs}, Batch {batch_idx}, Loss: {loss.item():.4f}')

    return train_loss


train_loss = train(net, train_loader, optimizer, criterion, device, epochs=5)
plot_curve(train_loss)


# Test the network
def test(net, test_loader, device):
    net.eval()
    total_correct = 0

    with torch.no_grad():
        for x, y in test_loader:
            x, y = x.to(device), y.to(device)
            output = net(x)
            pred = output.argmax(dim=1)
            total_correct += pred.eq(y).sum().item()

    accuracy = total_correct / len(test_loader.dataset)
    print(f'Test Accuracy: {accuracy:.4f}')


test(net, test_loader, device)


# Visualize some predictions
def visualize_predictions(net, test_loader, device):
    net.eval()
    x, y = next(iter(test_loader))
    x, y = x.to(device), y.to(device)
    output = net(x)
    pred = output.argmax(dim=1)
    plot_image(x.cpu(), pred.cpu(), 'test')


visualize_predictions(net, test_loader, device)
